import os
from typing import Union
from encryption_module import EncryptionModule


# 基础数据类型定义
BYTE = WORD = DWORD = int
BYTES = list[int]

# 消息类型定义
{%- for 一个消息类型 in 消息类型表 %}
MsgType{{一个消息类型[0]}} = {{一个消息类型[1]}}
"""{{一个消息类型[2]}}"""
{%- endfor %}

# 加密类型定义
EncModeNone = 0
"""不加密"""
EncModeSymmetric = 1
"""对称加密"""
EncModeAsymmetric = 2
"""非对称加密"""


# 私有函数定义
def _xor(frame: BYTES) -> BYTE:
    result = 0
    for num in frame:
        result ^= num
    return result


def _peek_word(frame: BYTES):
    word_value = 0
    word_value |= frame[0] << 8
    word_value |= frame[1]
    return word_value


def _push_byte(frame: BYTES, byte_value: BYTE):
    frame.append(byte_value)


def _push_word(frame: BYTES, word_value: WORD):
    frame.append((word_value >> 8) & 0xff)
    frame.append(word_value & 0xff)


def _push_dword(frame: BYTES, dword_value: DWORD):
    frame.append((dword_value >> 24) & 0xff)
    frame.append((dword_value >> 16) & 0xff)
    frame.append((dword_value >> 8) & 0xff)
    frame.append(dword_value & 0xff)


def _push_bytes(frame: BYTES, bytes_value: BYTES):
    frame.extend(bytes_value)


def _pop_byte(frame: BYTES) -> BYTE:
    byte_value = 0
    byte_value |= frame.pop(0)
    return byte_value


def _pop_word(frame: BYTES) -> WORD:
    word_value = 0
    word_value |= frame.pop(0) << 8
    word_value |= frame.pop(0)
    return word_value


def _pop_dword(frame: BYTES) -> DWORD:
    dword_value = 0
    dword_value |= frame.pop(0) << 24
    dword_value |= frame.pop(0) << 16
    dword_value |= frame.pop(0) << 8
    dword_value |= frame.pop(0)
    return dword_value


def _pop_bytes(frame: BYTES, bytes_len: int) -> BYTES:
    bytes_value = frame[:bytes_len].copy()
    del frame[:bytes_len]
    return bytes_value


class _MsgHead:
    """网络帧基类，包含头信息"""

    def __init__(self):
{%- for 消息结构表中的一个字段 in 消息结构表 %}
    {%- if [消息结构表中的一个字段[0]] == ['MsgHead'] %}
        self._{{消息结构表中的一个字段[0]}}: Union[{{消息结构表中的一个字段[2]}}, None] = 0x2323
    {%- elif [消息结构表中的一个字段[0]] == ['EncryptMode'] %}
        self._{{消息结构表中的一个字段[0]}}: Union[{{消息结构表中的一个字段[2]}}, None] = EncModeNone  # 默认不加密
    {%- else %}
        self._{{消息结构表中的一个字段[0]}}: Union[{{消息结构表中的一个字段[2]}}, None] = None
    {%- endif %}
{%- endfor %}
{%- for 消息结构表中的一个字段 in 消息结构表 %}

    def {{消息结构表中的一个字段[0]}}(self):
        """{{消息结构表中的一个字段[1]}}"""
        return self._{{消息结构表中的一个字段[0]}}
{%- endfor %}

    def SetEncryptMode(self, enc_mode: BYTE):
        """设置加密方式"""
        self._EncryptMode = enc_mode


{%- for 一个消息类型 in 消息类型表 %}


class Msg{{一个消息类型[0]}}(_MsgHead):
    """{{一个消息类型[2]}}"""
    {%- for 消息体中的一个字段 in 消息体字典[一个消息类型[0] + ' ' + 一个消息类型[1]] %}
    {{消息体中的一个字段[0]}}: {{消息体中的一个字段[2]}} = None
    """{{消息体中的一个字段[4]}}"""
    {%- endfor %}
    def __init__(self, uid: DWORD):
        super().__init__()
        self._UID = uid
        self._MsgType = MsgType{{一个消息类型[0]}}
{%- endfor %}


class NetProtocol:
    def __init__(self):
        """网络协议栈"""
        self.enc_module = EncryptionModule()

    def init_symmetric(self, symmetric_key: bytes, vi: bytes = b'0000000000000000'):
        """
        初始化协议栈的对称加密功能
        :param symmetric_key: 对称密钥
        :param vi: 初始向量
        """
        self.enc_module.init_symmetric(symmetric_key, vi)

    def init_asymmetric(self, remote_cert_path: str, curr_cert_path: str, auto_generate: bool = False) -> bool:
        """
        初始化协议栈的非对称加密功能
        :param remote_cert_path: 远端公钥证书路径
        :param curr_cert_path: 本地密钥证书路径
        :param auto_generate: 当curr_cert_path不存在时，是否主动生成密钥并导出到curr_cert_path
        """
        if os.path.exists(remote_cert_path):
            self.enc_module.import_asymmetric_remote_public_key(remote_cert_path)
        else:
            return False
        if os.path.exists(curr_cert_path):
            self.enc_module.import_asymmetric_local_private_key(curr_cert_path)
        else:
            if auto_generate:
                self.enc_module.generate_asymmetric_key()
                with open(curr_cert_path, mode='wb') as fp:
                    fp.write(self.enc_module.export_asymmetric_local_private_key())
                    fp.close()
                self.enc_module.export_asymmetric_local_private_key()
            else:
                return False
        return True

    def encode_msg(self, msg: Union[{%- for 一个消息类型 in 消息类型表 %}Msg{{一个消息类型[0]}},{{' '}}{%- endfor %}None]) -> Union[bytes, None]:
        """将消息类编码成可发送的网络帧"""
        # 编码消息体
        msg._MsgBody = []
{%- for idx in range(len(消息类型表)) %}
    {%- if [idx] == [0] %}
        if type(msg) is Msg{{消息类型表[idx][0]}}:
    {%- else %}
        elif type(msg) is Msg{{消息类型表[idx][0]}}:
    {%- endif %}
    {%- for 消息体中的一个字段 in 消息体字典[消息类型表[idx][0] + ' ' + 消息类型表[idx][1]] %}
            _push_{{消息体中的一个字段[2].lower()}}(msg._MsgBody, msg.{{消息体中的一个字段[0]}})
    {%- endfor %}
{%- endfor %}
        else:
            return None
        msg._MsgBodyLen = len(msg._MsgBody)
        msg_body_with_len = []
        _push_word(msg_body_with_len, msg._MsgBodyLen)
        _push_bytes(msg_body_with_len, msg._MsgBody)
        if msg._EncryptMode is EncModeNone:  # 不加密
            msg_body_cipher = msg_body_with_len
        elif msg._EncryptMode is EncModeSymmetric:  # 对称加密
            msg_body_cipher = list(self.enc_module.encrypt_symmetric(bytes(msg_body_with_len)))
        elif msg._EncryptMode is EncModeAsymmetric:  # 非对称加密
            signature = list(self.enc_module.sign_asymmetric(bytes(msg_body_with_len)))
            msg_body_cipher = list(self.enc_module.encrypt_asymmetric_remote(bytes(msg_body_with_len + signature)))
        else:
            return None
        if msg_body_cipher is None:
            return None
        # 编码消息头
        frame: BYTES = []
{%- for 消息结构表中的一个字段 in 消息结构表 %}
    {%- if [消息结构表中的一个字段[0]] == ['MsgLen'] %}
        msg._MsgLen = len(msg_body_cipher) + {{获取消息头长度(消息结构表)}}
        _push_{{消息结构表中的一个字段[2].lower()}}(frame, msg._MsgLen)
    {%- elif [消息结构表中的一个字段[0]] == ['CheckSum'] %}
        msg._CheckSum = _xor(frame[2:])
        _push_byte(frame, msg._CheckSum)
    {%- elif [消息结构表中的一个字段[0]] == ['MsgBodyLen'] %}
        _push_bytes(frame, msg_body_cipher)
    {%- elif [消息结构表中的一个字段[0]] == ['MsgBody'] %}
    {%- else %}
        _push_{{消息结构表中的一个字段[2].lower()}}(frame, msg._{{消息结构表中的一个字段[0]}})
    {%- endif %}
{%- endfor %}

        return bytes(frame)

    def decode_msg(self, frame: bytes) -> Union[{%- for 一个消息类型 in 消息类型表 %}Msg{{一个消息类型[0]}},{{' '}}{%- endfor %}None]:
        """将网络帧解码成易使用的消息类"""
        frame_arr = list(frame)
        # 解码消息头
        head = _MsgHead()
{%- for 消息结构表中的一个字段 in 消息结构表 %}
    {%- if [消息结构表中的一个字段[0]] == ['MsgBodyLen'] %}
        msg_body_cipher = _pop_bytes(frame_arr, len(frame) - {{获取消息头长度(消息结构表)}})
        if head._EncryptMode is EncModeNone:  # 不加密
            msg_body_with_len = msg_body_cipher
        elif head._EncryptMode is EncModeSymmetric:  # 对称加密
            msg_body_with_len = list(self.enc_module.decrypt_symmetric(bytes(msg_body_cipher)))
        elif head._EncryptMode is EncModeAsymmetric:  # 非对称加密
            msg_body_with_len_sign = list(self.enc_module.decrypt_asymmetric(bytes(msg_body_cipher)))
            if msg_body_with_len_sign is None:
                return None
            msg_body_len = _peek_word(msg_body_with_len_sign)
            if self.enc_module.verify_sign_asymmetric_remote(bytes(msg_body_with_len_sign[:msg_body_len+2]), bytes(msg_body_with_len_sign[msg_body_len+2:])):
                msg_body_with_len = msg_body_with_len_sign[:msg_body_len+2]
            else:
                msg_body_with_len = None
        else:
            return None
        if msg_body_with_len is None:
            return None
        head._MsgBodyLen = _pop_word(msg_body_with_len)
        head._MsgBody = msg_body_with_len
    {%- elif [消息结构表中的一个字段[0]] == ['MsgBody'] %}
    {%- else %}
        head._{{消息结构表中的一个字段[0]}} = _pop_{{消息结构表中的一个字段[2].lower()}}(frame_arr)
    {%- endif %}
{%- endfor %}
{%- for idx in range(len(消息类型表)) %}
    {%- if [idx] == [0] %}
        if head._MsgType == MsgType{{消息类型表[idx][0]}}:
    {%- else %}
        elif head._MsgType == MsgType{{消息类型表[idx][0]}}:
    {%- endif %}
            msg = Msg{{消息类型表[idx][0]}}(head.UID())
{%- endfor %}
        else:
            return None
{%- for 消息结构表中的一个字段 in 消息结构表 %}
        msg._{{消息结构表中的一个字段[0]}} = head._{{消息结构表中的一个字段[0]}}
{%- endfor %}
        # 解码消息体
        msg_body_buf = msg._MsgBody.copy()
{%- for idx in range(len(消息类型表)) %}
    {%- if [idx] == [0] %}
        if type(msg) is Msg{{消息类型表[idx][0]}}:
    {%- else %}
        elif type(msg) is Msg{{消息类型表[idx][0]}}:
    {%- endif %}
    {%- for 消息体中的一个字段 in 消息体字典[消息类型表[idx][0] + ' ' + 消息类型表[idx][1]] %}
        {%- if [消息体中的一个字段[2].lower()] == ["bytes"] %}
            msg.{{消息体中的一个字段[0]}} = _pop_{{消息体中的一个字段[2].lower()}}(msg_body_buf, msg.{{消息体中的一个字段[3]}})
        {%- else %}
            msg.{{消息体中的一个字段[0]}} = _pop_{{消息体中的一个字段[2].lower()}}(msg_body_buf)
        {%- endif %}
    {%- endfor %}
{%- endfor %}
        else:
            return None
        return msg

    @staticmethod
    def decode_msg_head(frame: bytes) -> Union[_MsgHead, None]:
        """从网络帧解析消息头，消息体不会被解析"""
        frame_arr = list(frame[: {{获取消息头长度(消息结构表) - 1}}])
        # 解码消息头
        head = _MsgHead()
{%- for 消息结构表中的一个字段 in 消息结构表 %}
    {%- if [消息结构表中的一个字段[0]] == ['MsgBodyLen'] or [消息结构表中的一个字段[0]] == ['MsgBody'] %}
    {%- elif [消息结构表中的一个字段[0]] == ['CheckSum'] %}
        head._{{消息结构表中的一个字段[0]}} = frame[-1]
    {%- else %}
        head._{{消息结构表中的一个字段[0]}} = _pop_{{消息结构表中的一个字段[2].lower()}}(frame_arr)
    {%- endif %}
{%- endfor %}

        return head

    @staticmethod
    def cut_frame(input_buf: BYTES) -> Union[bytes, None]:
        """将输入缓冲区中的数据切包成一个网络帧并返回，无效数据将被自动舍弃"""
        # 寻找消息头
        msg_start_idx = -1
        for idx in range(0, len(input_buf) - 1):
            if (input_buf[idx] << 8) | input_buf[idx + 1] == 0x2323:
                msg_start_idx = idx
                break
        if msg_start_idx == -1:
            input_buf.clear()
            return None
        if msg_start_idx > 0:
            del input_buf[: msg_start_idx]
        if len(input_buf) < 4:
            return None
        msg_len = (input_buf[2] << 8) | input_buf[3]
        if len(input_buf) < msg_len:
            return None
        if _xor(input_buf[: msg_len - 1]) == input_buf[msg_len - 1]:
            frame = bytes(input_buf[:msg_len])
        else:
            frame = None
        if len(input_buf) == msg_len:
            input_buf.clear()
        else:
            del input_buf[: msg_len]

        return frame
